package com.sluggard.ai.controller;

import com.sluggard.mp.model.req.IdReq;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.sluggard.ai.config.MilvusProperties;
import com.sluggard.ai.manage.VectorStoreManage;
import com.sluggard.ai.model.entity.KnowledgeEntity;
import com.sluggard.ai.model.req.KnowledgeReq;
import com.sluggard.ai.service.IKnowledgeService;
import com.sluggard.mp.model.req.IdsReq;
import com.sluggard.mp.mvc.BaseController;

import io.milvus.client.MilvusClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.DescribeDatabaseResponse;
import io.milvus.grpc.GetLoadStateResponse;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CollectionSchemaParam;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.CreateDatabaseParam;
import io.milvus.param.collection.DescribeCollectionParam;
import io.milvus.param.collection.DescribeDatabaseParam;
import io.milvus.param.collection.DropCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.GetLoadStateParam;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;

/**
 * <p>
 * ai 知识库 前端控制器
 * </p>
 *
 * @author lizheng
 * @since 2025-04-23
 */
@Slf4j
@Tag(name = "ai 知识库")
@RestController
@RequestMapping("/knowledge")
public class KnowledgeController extends BaseController<KnowledgeReq, KnowledgeEntity, IKnowledgeService> {

    @Resource
    private MilvusClient milvusClient;

    @Resource
    private MilvusProperties milvusProperties;

    @Override
    protected void afterSave(KnowledgeReq req, KnowledgeEntity instance) {
        // 知识库创建好了之后,在向量数据库中创建一个集合
        // 判断database 是否存在
        R<DescribeDatabaseResponse> describeDatabase = milvusClient
                .describeDatabase(DescribeDatabaseParam.newBuilder()
                        .withDatabaseName(milvusProperties.getDatabaseName())
                        .build());
        if (describeDatabase.getStatus() != 0) {
            CreateDatabaseParam createDatabaseParam = CreateDatabaseParam.newBuilder()
                    .withDatabaseName(milvusProperties.getDatabaseName())
                    .build();
            R<RpcStatus> database = milvusClient.createDatabase(createDatabaseParam);
            if (database.getStatus() != 0) {
                log.error("创建数据库 {} 失败 ====> 状态码：{} 原因：{}", milvusProperties.getDatabaseName(),
                        database.getStatus(), database.getMessage());
                throw new RuntimeException("创建集合失败");
            } else {
                log.info("创建数据库 {} ", milvusProperties.getDatabaseName());
            }
        }

        // 判断集合是否存在
        Integer status1 = milvusClient.describeCollection(DescribeCollectionParam.newBuilder()
                .withDatabaseName(milvusProperties.getDatabaseName())
                .withCollectionName(VectorStoreManage.getCollectionName(instance.getId()))
                .build()).getStatus();
        if (status1 != 0) {
            CollectionSchemaParam collectionSchemaParam = CollectionSchemaParam.newBuilder()
                    .addFieldType(FieldType.newBuilder()
                            .withAutoID(true)
                            .withDataType(DataType.Int64)
                            .withPrimaryKey(true)
                            .withName("id")
                            .build())
                    .addFieldType(FieldType.newBuilder()
                            .withDataType(DataType.VarChar)
                            .withMaxLength(100)
                            .withName("doc_id")
                            .build())
                    .addFieldType(FieldType.newBuilder()
                            .withDimension(milvusProperties.getDimension())
                            .withDataType(DataType.FloatVector)
                            .withName("embedding")
                            .build())
                    .addFieldType(FieldType.newBuilder()
                            .withDimension(milvusProperties.getDimension())
                            .withMaxLength(1024)
                            .withDataType(DataType.VarChar)
                            .withName("content")
                            .build())
                    .addFieldType(FieldType.newBuilder()
                            .withDimension(milvusProperties.getDimension())
                            .withDataType(DataType.JSON)
                            .withName("metadata")
                            .build())
                    .build();

            CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                    .withSchema(collectionSchemaParam)
                    .withDatabaseName(milvusProperties.getDatabaseName())
                    .withCollectionName(VectorStoreManage.getCollectionName(instance.getId()))
                    .withShardsNum(2)
                    .build();
            R<RpcStatus> collection = milvusClient.createCollection(createCollectionParam);
            if (collection.getStatus() != 0) {
                log.error("创建数据库 {} 集合 {} 失败 ====> 状态码：{} 原因：{}", milvusProperties.getDatabaseName(),
                        VectorStoreManage.getCollectionName(instance.getId()),
                        collection.getStatus(), collection.getMessage());
                throw new RuntimeException("创建集合失败");
            } else {
                R<GetLoadStateResponse> loadState = milvusClient.getLoadState(GetLoadStateParam
                        .newBuilder()
                        .withDatabaseName(milvusProperties.getDatabaseName())
                        .withCollectionName(
                                VectorStoreManage.getCollectionName(instance.getId()))
                        .build());
                if (loadState.getStatus() == 0) {
                    log.info("创建数据库 {} 集合 {} 成功", milvusProperties.getDatabaseName(),
                            VectorStoreManage.getCollectionName(instance.getId()));
                } else {
                    log.error("加载数据库 {} 集合 {} 失败 ====> 状态码：{} 原因：{}",
                            milvusProperties.getDatabaseName(),
                            VectorStoreManage.getCollectionName(instance.getId()),
                            loadState.getStatus(), loadState.getMessage());
                }

            }
        }
    }

    @Override
    protected void onDelete(IdsReq idsReq) {
        idsReq.getIds().parallelStream()
                .forEach(item -> milvusClient.dropCollection(DropCollectionParam.newBuilder()
                        .withDatabaseName(milvusProperties.getDatabaseName())
                        .withCollectionName(VectorStoreManage.getCollectionName(item))
                        .build()));
    }

    @Override
    protected void onDeleteById(IdReq idReq) {
        milvusClient.dropCollection(DropCollectionParam.newBuilder()
                .withDatabaseName(milvusProperties.getDatabaseName())
                .withCollectionName(VectorStoreManage.getCollectionName(idReq.getId()))
                .build());
    }
}
